from typing import Dict, List
from .node_handler import NodeHandler,FollowHandler
from ..generator import ScaleDotProductGenerator, StrategyGenerator
from ..StrategiesVector import StrategiesVector
from .registry import operator_registry
from ..shard.placement_types import DeviceMesh,AbstractTensor,OperationData, OperationDataType
from geesibling.core.types import Graph, Node


__all__ = ['ScaleDotProductHandler']


@operator_registry.register('aten._scaled_dot_product_efficient_attention.default') # 不确定
class ScaleDotProductHandler(FollowHandler):
    """
    A UnaryElementwiseHandler which deals with the sharding strategies for Elementwise Op.
    """

    def get_strategy_generator(self) -> List[StrategyGenerator]:
        op_data_mapping = self.get_operation_data_mapping()
        generators = []
        generators.append(ScaleDotProductGenerator(op_data_mapping, self.device_mesh, self.pre_strategies_vectors))
        return generators

    def get_operation_data_mapping(self) -> Dict[str, OperationData]:
        # 一个输入一个输出
        # 得到node的输入node
        mapping = {}
        for i in range(len(self.node.inputs)):
            if self.node.input_shape(i):
                input_abs_data = AbstractTensor(self.node.input_shape(i),self.node.input_type(i))
                physical_input = OperationData(name=str(self.node.inputs[i]), type=OperationDataType.ARG, data=input_abs_data)
                mapping["input_{}".format(i)] = physical_input
        output_abs_data = AbstractTensor(self.node.output_shape(0),self.node.output_type(0))
        physical_output = OperationData(name=str(self.node), data=output_abs_data,type=OperationDataType.OUTPUT)
        mapping['output'] = physical_output
        return mapping
